# 导入必要的库
import torch  # PyTorch核心库
import torch.nn as nn  # PyTorch神经网络库
import torch.optim as optim  # PyTorch优化器库
import torchvision  # PyTorch的图像处理库
import torchvision.transforms as transforms  # 图像转换工具

# 定义超参数
batch_size = 64  # 每个训练批次中的图像数量
learning_rate = 0.001  # 学习率，控制模型参数更新的速度
num_epochs = 5  # 训练周期数，即遍历整个训练集的次数

# 加载MNIST数据集
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])  # 图像转换流程：转为张量并归一化
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transform, download=True)  # 加载训练数据集
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, transform=transform)  # 加载测试数据集

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)  # 创建训练数据加载器
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)  # 创建测试数据加载器

# 定义神经网络模型
class SimpleNN(nn.Module):  # 定义一个名为SimpleNN的神经网络类，继承自nn.Module
    def __init__(self):
        super(SimpleNN, self).__init__()  # 调用父类的初始化方法
        self.fc1 = nn.Linear(28 * 28, 512)  # 定义一个全连接层，输入特征数为28*28（图像展平后的维度），输出特征数为512
        self.fc2 = nn.Linear(512, 256)  # 第二个全连接层，输入特征数为512，输出特征数为256
        self.fc3 = nn.Linear(256, 10)  # 第三个全连接层，输入特征数为256，输出特征数为10（对应10个类别）

    def forward(self, x):  # 定义前向传播方法
        x = x.view(-1, 28 * 28)  # 将输入图像展平为一维向量，-1表示自动计算该维度的大小
        x = torch.relu(self.fc1(x))  # 通过第一个全连接层，并使用ReLU激活函数
        x = torch.relu(self.fc2(x))  # 通过第二个全连接层，并使用ReLU激活函数
        # 注意：CrossEntropyLoss内部已经包含了softmax，所以这里不需要再对输出使用softmax
        x = torch.softmax(self.fc3(x), dim=1)  # 通过第三个全连接层，并对输出使用softmax激活函数（实际上在CrossEntropyLoss中不需要这一步）
        return x

# 初始化模型、损失函数和优化器
model = SimpleNN()  # 实例化SimpleNN模型
criterion = nn.CrossEntropyLoss()  # 定义损失函数为交叉熵损失
optimizer = optim.Adam(model.parameters(), lr=learning_rate)  # 定义优化器为Adam，并传入模型参数和学习率

# 训练模型
for epoch in range(num_epochs):  # 遍历训练周期
    for i, (images, labels) in enumerate(train_loader):  # 遍历训练数据加载器
        # 前向传播
        outputs = model(images)  # 将图像输入模型，得到预测输出
        loss = criterion(outputs, labels)  # 计算损失

        # 反向传播和优化
        optimizer.zero_grad()  # 清空梯度
        loss.backward()  # 反向传播计算梯度
        optimizer.step()  # 更新模型参数

        if (i + 1) % 100 == 0:  # 每100个批次打印一次损失
            print(f'Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{len(train_loader)}], Loss: {loss.item():.4f}')

# 评估模型
model.eval()  # 将模型设置为评估模式，这会关闭Dropout等层的训练特性
with torch.no_grad():  # 禁用梯度计算，以加速评估并减少内存消耗
    correct = 0  # 初始化正确预测的数量
    total = 0  # 初始化总预测的数量
    for images, labels in test_loader:  # 遍历测试数据加载器
        outputs = model(images)  # 将图像输入模型，得到预测输出
        _, predicted = torch.max(outputs.data, 1)  # 在输出上应用argmax得到预测的类别
        total += labels.size(0)  # 更新总预测的数量
        correct += (predicted == labels).sum().item()  # 计算并更新正确预测的数量

    print(f'Accuracy of the model on the test images: {100 * correct / total:.2f}%')  # 打印测试准确率

# 保存模型
torch.save(model.state_dict(), 'simple_nn.pth')  # 将模型的状态字典保存到文件